package heap

import "container/heap"

func topKFrequent(nums []int, k int) []int {
	cnt := map[int]int{}
	for _, num := range nums {
		cnt[num]++
	}
	h := &tHeap{}
	heap.Init(h)
	for k, v := range cnt {
		heap.Push(h, [2]int{k, v})
		if h.Len() > k {
			heap.Pop(h)
		}
	}
	res := make([]int, k)
	for i := 0; i < k; i++ {
		res[k-i-1] = heap.Pop(h).([2]int)[0]
	}
	return res
}

type tHeap [][2]int

// Len is the number of elements in the collection.
func (t tHeap) Len() int {
	return len(t)
}

func (t tHeap) Less(i int, j int) bool {
	return t[i][1] < t[j][1]
}

// Swap swaps the elements with indexes i and heap
func (t tHeap) Swap(i int, j int) {
	t[i], t[j] = t[j], t[i]
}

func (t *tHeap) Push(x any) {
	*t = append(*t, x.([2]int))
}

func (t *tHeap) Pop() any {
	old := *t
	n := len(old)
	x := old[n-1]
	*t = old[0 : n-1]
	return x
}
